
Program lin_inversion

use metro_mod           ! Module contains invert function
use netcdf

Implicit none

!  Purpose:
!   Simple Bayesian inversion based on given measurements, uncertainties and sensitivites
!  
!  Inputs: (statesize,nMeasure, y, y_error, x_ap, H, x_ap_sd) in .nc file
!   
!  Outputs:
!   ./outputs/output_file.nc contains optimized state vector (x) and optimized observations (y_opt) 
!
!  Requires:
!   Netcdf file *input_file.nc*, stored in ./inputs
!   
!   Netcdf & LAPACK libraries
!   
!  Modules Used: metro_mod.f90
!   
!  History:
! 	Written by: Somebody, UoB, Dec 13, 2013
!

 integer                           :: statesize, nmeasure 
 double precision                  :: x_ap_sd   
 integer                           :: xi, yi, status                        ! counters
 double precision,allocatable, dimension(:)    :: y, y_error, x, x_ap, y_opt                                      
 double precision,allocatable, dimension(:,:)  :: H, P1, R1, P, a,b           
 integer                           :: errstat                               ! for error checking
 integer                           :: inputID, outputID, &                  ! Required for Netcdf files and variables
                                       statesizeID, nMeasureID, &
                                       x_apID, yID, y_errorID, HID, x_ap_sdID, &
                                       statesizeDID, nMeasureDID, xID, y_optID  

! ---------------- Open Netcdf input and output files --------------------
 STATUS = NF90_OPEN('./inputs/input_file.nc', NF90_NOWRITE, ncid=inputID)
 STATUS = NF90_CREATE ('./outputs/output_file.nc', cmode=0, ncid=outputID)

 STATUS = NF90_INQ_VARID(inputID, 'statesize', statesizeID)
 STATUS = NF90_INQ_VARID(inputID, 'nmeasure', nMeasureID)

 STATUS = NF90_GET_VAR(inputID, statesizeID, statesize)
 STATUS = NF90_GET_VAR(inputID, nMeasureID, nMeasure)

! ------------ Allocate space for all variables ---------------------
allocate(x_ap(statesize),stat=errstat)                 ! Allocate x_ap
  if (errstat /= 0) then
     write (*,*) 'ERROR: could not allocate x_ap array'
     stop
  endif

allocate(x(statesize),stat=errstat)                    ! Allocate x
  if (errstat /= 0) then
     write (*,*) 'ERROR: could not allocate x array'
     stop
  endif

allocate(P1(statesize,statesize),stat=errstat)         ! Allocate P1
  if (errstat /= 0) then
     write (*,*) 'ERROR: could not allocate P1 array'
     stop
  endif 


allocate(y(nMeasure),stat=errstat)                     ! Allocate y
  if (errstat /= 0) then
     write (*,*) 'ERROR: could not allocate y array'
     stop
  endif

allocate(y_error(nMeasure),stat=errstat)               ! Allocate y_error
  if (errstat /= 0) then
     write (*,*) 'ERROR: could not allocate y_error array'
     stop
  endif

allocate(R1(nMeasure,nMeasure),stat=errstat)           ! Allocate R1
  if (errstat /= 0) then
     write (*,*) 'ERROR: could not allocate R1 array'
     stop
  endif


allocate(H(nMeasure,statesize),stat=errstat)           ! Allocate H
  if (errstat /= 0) then
     write (*,*) 'ERROR: could not allocate H array'
     stop
  endif

allocate(P(statesize,statesize),stat=errstat)           ! Allocate P
  if (errstat /= 0) then
     write (*,*) 'ERROR: could not allocate P array'
     stop
  endif

allocate(y_opt(nMeasure),stat=errstat)                  ! Allocate y_opt
  if (errstat /= 0) then
     write (*,*) 'ERROR: could not allocate y_opt array'
     stop
  endif

! -------------- Read in variables from input netcdf ---------------------
   STATUS = NF90_INQ_VARID (inputID, 'x_ap', x_apID)
   STATUS = NF90_INQ_VARID (inputID, 'y', yID)
   STATUS = NF90_INQ_VARID (inputID, 'y_error', y_errorID)
   STATUS = NF90_INQ_VARID (inputID, 'H', HID)
   STATUS = NF90_INQ_VARID (inputID, 'x_ap_sd', x_ap_sdID)
   
   STATUS = NF90_GET_VAR (inputID, x_apID, x_ap)
   STATUS = NF90_GET_VAR (inputID, yID, y)
   STATUS = NF90_GET_VAR (inputID, y_errorID, y_error)
   STATUS = NF90_GET_VAR (inputID, HID, H)
   STATUS = NF90_GET_VAR (inputID, x_ap_sdID, x_ap_sd)

   STATUS = nf90_close(inputID)

! ------------- Calculate diagonal covariance matrices ------------------
  Do xi=1,statesize
     P1(xi,xi) = (1./(x_ap(xi)*x_ap_sd))**2
  Enddo
  Do yi=1,nMeasure
     R1(yi,yi)=1./(y_error(yi)**2)
  Enddo

! -------------------- Linear Inversion --------------------
  Write(*,*) 'Starting linear inversion...'

  P = Ainv(matmul(matmul(transpose(H),R1),H) + P1)               ! Function Ainv contained in metro_mod
  x = x_ap + matmul(matmul(matmul(P,transpose(H)),R1),(y-matmul(H,x_ap))) 
  y_opt = matmul(H,x)

  Write(*,*) '...Finished linear inversion.'

! ----------------- Write to Netcdf -------------------------------
  STATUS = NF90_DEF_DIM (outputID, 'statesize', statesize, statesizeDID)
  STATUS = NF90_DEF_DIM (outputID, 'nmeasure', nMeasure, nMeasureDID)

  STATUS = NF90_DEF_VAR (outputID, 'x', NF90_DOUBLE, statesizeDID, xID)
  STATUS = NF90_DEF_VAR (outputID, 'y_opt', NF90_DOUBLE, nMeasureDID, y_optID)

  STATUS = NF90_ENDDEF (outputID)

  STATUS = NF90_PUT_VAR (outputID, xID, x)
  STATUS = NF90_PUT_VAR (outputID, yID, y_opt)

  STATUS = NF90_CLOSE(outputID)
  
! ------------------ Deallocate arrays --------------------------
  if(allocated(x))   deallocate(x,stat=errstat)         
  if(allocated(x_ap))   deallocate(x_ap,stat=errstat)               
  if(allocated(y))   deallocate(y,stat=errstat)
  if(allocated(y_error))   deallocate(y_error,stat=errstat)
  if(allocated(y_opt))   deallocate(y_opt,stat=errstat)
  if(allocated(R1))   deallocate(R1,stat=errstat)
  if(allocated(H))   deallocate(H,stat=errstat)
  if(allocated(P1))   deallocate(P1,stat=errstat)
  if(allocated(P))   deallocate(P,stat=errstat)

End program lin_inversion


